# coding=utf-8
### 消费者
import json
import time
import pika
import os
import threading

from utils.oss_file import download_file_to_temp_folder
from omegaconf import OmegaConf
from dataset.data_factory import classify_provider, provider, unsupervised_provider
from train_classify import TrainVal as ClassifyTrainVal
from train_segment import TrainVal as SegmentTrainVal
from train_unsupervised import TrainVal as UnsupervisedTrainVal
import socket
import torch
import functools

"""                       
rabbitmq消息体：
body = {
    'type': 'classify',
    'train_task_id': 1,
    'yaml_path': 'http://msjava.oss-cn-qingdao.aliyuncs.com/classify_config.yaml'
}
"""


def call_for_train_classify(config):
    train_loader, valid_loader = classify_provider(
        config.dataset.type,
        config.dataset.data_folder,
        config.batch_size,
        config.num_workers,
        config.dataset.height,
        config.dataset.width,
        config.dataset.mean,
        config.dataset.std,
        config.augmentation_flag
    )
    train_val = ClassifyTrainVal(config, device)
    train_val.train(train_loader, valid_loader)


"""
rabbitmq消息体：
body = {
    'type': 'segment',
    'train_task_id': 1,
    'yaml_path': 'http://msjava.oss-cn-qingdao.aliyuncs.com/classify_config.yaml'
}
"""


def call_for_train_segment(config):
    train_loader, valid_loader = provider(
        config.dataset.type,
        config.dataset.data_folder,
        config.batch_size,
        config.num_workers,
        config.dataset.height,
        config.dataset.width,
        config.dataset.mean,
        config.dataset.std,
        config.augmentation_flag,
        config.mask_only
    )
    train_val = SegmentTrainVal(config, device)
    train_val.train(train_loader, valid_loader)


"""
rabbitmq消息体：
body = {
    'type': 'unsupervised',
    'train_task_id': 3,
    'yaml_path': 'http://msjava.oss-cn-qingdao.aliyuncs.com/unsupervised_config.yaml'
}
"""


def call_for_train_unsupervised(config):
    train_loader, valid_loader = unsupervised_provider(
        config.dataset.type,
        config.dataset.data_folder,
        config.batch_size,
        config.num_workers,
        config.dataset.height,
        config.dataset.width,
        config.dataset.mean,
        config.dataset.std,
        config.augmentation_flag
    )
    train_val = UnsupervisedTrainVal(config, device)
    train_val.train(train_loader, valid_loader)


def ack_message(channel, delivery_tag):
    if channel.is_open:
        channel.basic_ack(delivery_tag)
    else:
        # Channel is already closed, so we can't ACK this message;
        # log and/or do something that makes sense for your app in this case.
        pass


def on_message(channel, method_frame, header_frame, body):
    print(f'on_message thread id: {threading.get_ident()}')
    delivery_tag = method_frame.delivery_tag
    t = threading.Thread(target=do_work, args=(channel, delivery_tag, body))
    t.start()


# 回调函数
def do_work(channel, delivery_tag, body):
    body = OmegaConf.create(json.loads(body))
    yaml_path = body.yaml_path
    local_yaml_path = download_file_to_temp_folder(yaml_path, temp_path)
    config = OmegaConf.load(local_yaml_path)
    config = OmegaConf.merge(config, body)
    ###############################
    # config.epoch = 2
    # config.batch_size = 2
    ###############################
    os.remove(local_yaml_path)
    if config.type == 'classify':
        call_for_train_classify(config)
    elif config.type == 'segment':
        call_for_train_segment(config)
    elif config.type == 'unsupervised':
        call_for_train_unsupervised(config)

    cb = functools.partial(ack_message, channel, delivery_tag)
    channel.connection.add_callback_threadsafe(cb)


if __name__ == '__main__':
    temp_path = ''
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    hostname = socket.gethostname()
    ip_address = socket.gethostbyname(hostname)

    user_info = pika.PlainCredentials('admin', 'admin')
    # user_info = pika.PlainCredentials('xyf', '123456')

    connection = pika.BlockingConnection(
        pika.ConnectionParameters('192.168.1.111', 5672, virtual_host='/', credentials=user_info, heartbeat=600,
                                  blocked_connection_timeout=300))
    # connection = pika.BlockingConnection(
    #     pika.ConnectionParameters('119.91.47.43', 5672, virtual_host='/', credentials=user_info))
    channel = connection.channel()
    # queue_name = f'ip:{ip_address}'
    queue_name = 'ip:192.168.1.180'
    channel.queue_declare(queue=queue_name, durable=True)
    channel.basic_qos(prefetch_count=1)
    channel.basic_consume(queue=queue_name,  # 接收指定queue的消息
                          auto_ack=False,  # 指定为False，表示取消自动应答，交由回调函数手动应答
                          on_message_callback=on_message  # 设置收到消息的回调函数
                          )
    channel.start_consuming()
